# -*- coding: utf-8 -*-
'''
File:modeling_barranzi2LM.py                                         
Time:2025/9/29 15:43                                    
IDE:PyCharm                                     
Author:Barranzi An                                        
email:awc19930818@outlook.com                           
github:https://github.com/La0bALanG                     
Barranzi's Blog:私聊巴郎子要文档链接                  
requirement:(Please describle your requirement here) -->
'''

from typing import Optional, Tuple, Dict, Callable
from packaging import version

import torch
from torch import nn
import torch.nn.functional as F
import math

from .configuration_barranzi2 import Config


class RMSNorm(nn.Module):
    """
    均方根归一化（RMSNorm）层实现。

    RMSNorm是一种归一化技术，通过输入的均方根值对输入进行归一化处理。
    与LayerNorm不同，RMSNorm不通过减去均值来中心化输入，这使得它在计算上更高效，
    同时在许多transformer模型中保持相似的性能。

    归一化计算公式：output = (input / sqrt(mean(input^2) + eps)) * weight
    """

    def __init__(self, hidden_size, eps=1e-6):
        """
        初始化RMSNorm层。

        参数:
            hidden_size (int): 输入特征的维度
            eps (float, optional): 添加到分母的小常数，用于数值稳定性。
                                 默认值为1e-6。
        """
        super().__init__()
        # 可学习的权重参数，用于缩放归一化后的输出
        # 初始化为全1，允许模型学习最优的缩放比例
        self.weight = nn.Parameter(torch.ones(hidden_size))
        # 小的epsilon值，防止除零错误并确保数值稳定性
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        """
        RMSNorm的前向传播。

        参数:
            hidden_states (torch.Tensor): 形状为(..., hidden_size)的输入张量

        返回:
            torch.Tensor: 与输入形状相同的归一化张量
        """
        # 存储原始数据类型，以便在计算后转换回来
        input_dtype = hidden_states.dtype

        # 转换为float32以确保计算过程中的数值稳定性
        hidden_states = hidden_states.to(torch.float32)

        # 沿最后一个维度计算平方值的均值
        # 这给出了方差（没有中心化）
        # keepdim=True保持维度以便广播
        variance = hidden_states.pow(2).mean(-1, keepdim=True)

        # 通过除以方差加epsilon的平方根来归一化输入
        # torch.rsqrt(x)等价于1 / sqrt(x)，但数值上更稳定
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)

        # 应用可学习的权重并转换回原始数据类型
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        """
        提供模块的额外字符串表示。

        返回:
            str: 包含权重形状和epsilon值的字符串
        """
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class MLP(nn.Module):
    """
    多层感知机（MLP）模块，实现了SwiGLU激活函数。

    这是大型语言模型中常用的前馈网络结构，采用门控机制来增强模型的表达能力。
    SwiGLU通过将输入分成两个路径，一个路径经过激活函数作为门控信号，
    另一个路径直接通过，然后将两者相乘，最后通过下投影层输出。

    结构：input -> [gate_proj -> SiLU] * [up_proj] -> down_proj -> output
    """

    def __init__(self, config: Config, intermediate_size: Optional[int] = None):
        """
        初始化MLP模块。

        参数:
            config (Config): 模型配置对象，包含隐藏层大小等参数
            intermediate_size (Optional[int]): 中间层大小，如果为None则使用config中的值
        """
        super().__init__()
        # 确定中间层大小，优先使用传入的参数，否则使用配置中的默认值
        config_intermediate_size = intermediate_size if intermediate_size else config.intermediate_size

        # 门控投影层：将输入从hidden_size映射到intermediate_size
        # 用于生成门控信号，经过SiLU激活函数
        self.gate_proj = nn.Linear(config.hidden_size, config_intermediate_size, bias=False)

        # 上投影层：将输入从hidden_size映射到intermediate_size
        # 与门控信号相乘，实现特征的选择性增强
        self.up_proj = nn.Linear(config.hidden_size, config_intermediate_size, bias=False)

        # 下投影层：将intermediate_size映射回hidden_size
        # 将中间层的特征压缩回原始维度
        self.down_proj = nn.Linear(config_intermediate_size, config.hidden_size, bias=False)

        # 激活函数：使用SiLU（Sigmoid Linear Unit）函数
        # SiLU(x) = x * sigmoid(x)，具有平滑且非单调的特性
        self.activation = nn.SiLU()

    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        MLP的前向传播，实现SwiGLU激活函数。

        计算过程：
        1. 输入通过gate_proj和up_proj分别投影到中间维度
        2. gate_proj的输出经过SiLU激活函数作为门控信号
        3. 门控信号与up_proj的输出逐元素相乘
        4. 结果通过down_proj投影回原始维度

        参数:
            hidden_states (torch.Tensor): 输入张量，形状为(batch_size, seq_len, hidden_size)

        返回:
            Tuple[torch.Tensor, Optional[torch.Tensor]]:
                - 输出张量，形状与输入相同
                - 第二个元素为None（用于MoE模块的兼容性）
        """
        # 实现SwiGLU激活函数：down_proj(SiLU(gate_proj(x)) * up_proj(x))
        # 这种结构允许模型学习特征的非线性组合，增强表达能力
        return self.down_proj(self.activation(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))


def _compute_default_rope_parameters(
        config: Config,
        device: Optional["torch.device"] = None,
        seq_len: Optional[int] = None,
) -> Tuple["torch.Tensor", float]:
    """
    根据原始RoPE实现计算逆频率参数。

    RoPE（Rotary Position Embedding）是一种位置编码技术，通过旋转操作将位置信息
    嵌入到查询和键向量中，避免了传统位置编码的插值问题，同时保持了相对位置信息。

    参数:
        config ([`~transformers.PretrainedConfig`]):
            模型配置对象，包含RoPE相关的超参数
        device (`torch.device`):
            用于初始化逆频率张量的设备
        seq_len (`int`, *optional*):
            当前序列长度。对于默认类型的RoPE未使用

    返回:
        Tuple[`torch.Tensor`, `float`]: 包含RoPE嵌入的逆频率张量和后处理缩放因子的元组
            - inv_freq: 逆频率张量，用于生成旋转位置编码
            - attention_factor: 注意力缩放因子（在此类型RoPE中未使用，固定为1.0）
    """
    # 获取RoPE的基础频率参数，通常为10000
    # 这个参数控制旋转频率的基数值，影响位置编码的周期性
    base = config.rope_config.rope_theta

    # 获取部分旋转因子，决定应用RoPE的维度比例
    # 例如0.5表示只对一半的特征维度应用旋转
    partial_rotary_factor = config.rope_config.partial_rotary_factor

    # 计算注意力头的维度
    # 优先使用配置中的head_dim，否则通过hidden_size除以注意力头数计算
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)

    # 计算实际应用RoPE的维度
    # 通过头维度乘以部分旋转因子得到实际旋转的维度大小
    dim = int(head_dim * partial_rotary_factor)

    # 注意力缩放因子，在默认RoPE中未使用，固定为1.0
    attention_factor = 1.0  # Unused in this type of RoPE

    # 计算逆频率张量
    # 这是RoPE的核心计算，生成用于旋转位置编码的频率参数
    # 公式：inv_freq[i] = 1 / (base^(i/dim))，其中i为偶数索引
    # 这种设计使得不同维度具有不同的旋转频率，低频维度变化慢，高频维度变化快
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))

    return inv_freq, attention_factor


def _compute_dynamic_ntk_parameters(
        config: Config,
        device: Optional["torch.device"] = None,
        seq_len: Optional[int] = None,
) -> Tuple["torch.Tensor", float]:
    """
    使用NTK（Neural Tangent Kernel）缩放计算逆频率参数。

    动态NTK是一种改进的RoPE变体，能够更好地处理长序列外推问题。
    通过动态调整基础频率，使得模型在处理超出训练时长度的序列时，
    仍能保持较好的性能。该方法由Reddit用户/u/bloc97和/u/emozilla提出。

    参数:
        config ([`~transformers.PretrainedConfig`]):
            模型配置对象，包含RoPE相关的超参数
        device (`torch.device`):
            用于初始化逆频率张量的设备
        seq_len (`int`, *optional*):
            当前序列长度，用于在推理时动态更新RoPE参数

    返回:
        Tuple[`torch.Tensor`, `float`]: 包含RoPE嵌入的逆频率张量和后处理缩放因子的元组
            - inv_freq: 经过NTK缩放调整的逆频率张量
            - attention_factor: 注意力缩放因子（在此类型RoPE中未使用，固定为1.0）
    """
    # 获取RoPE的基础频率参数，通常为10000
    # 这个参数会在后续根据序列长度进行动态调整
    base = config.rope_config.rope_theta

    # 获取部分旋转因子，决定应用RoPE的维度比例
    partial_rotary_factor = config.rope_config.partial_rotary_factor

    # 计算注意力头的维度
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)

    # 计算实际应用RoPE的维度
    dim = int(head_dim * partial_rotary_factor)

    # 获取模型支持的最大位置嵌入长度
    max_position_embeddings = config.max_position_embeddings

    # 获取NTK缩放因子，用于控制外推程度
    factor = config.rope_config.factor

    # 注意力缩放因子，在NTK RoPE中未使用，固定为1.0
    attention_factor = 1.0  # Unused in this type of RoPE

    # 处理序列长度参数
    # 如果seq_len为None或不超过最大位置嵌入长度，则使用最大位置嵌入长度
    # 这确保了在初始化或短序列情况下使用默认参数
    seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings

    # 动态调整基础频率以适应长序列外推
    # 这是NTK方法的核心公式，通过调整base来改善长序列处理能力
    # 公式：base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))
    # 其中：
    # - (factor * seq_len / max_position_embeddings) - (factor - 1) 计算缩放比例
    # - (dim / (dim - 2)) 是维度相关的指数调整
    base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))

    # 使用调整后的基础频率计算逆频率张量
    # 公式：inv_freq[i] = 1 / (base^(i/dim))，其中i为偶数索引
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))

    return inv_freq, attention_factor


def _compute_yarn_parameters(
        config: Config,
        device: Optional["torch.device"] = None,
        seq_len: Optional[int] = None
) -> Tuple["torch.Tensor", float]:
    """
    使用YaRN (Yet another RoPE extensioN)方法计算RoPE的逆频率和注意力缩放因子。
    YaRN是一种改进的RoPE扩展技术，能够有效处理长序列外推问题，通过调整频率参数和注意力缩放来提升模型在长上下文中的性能。
    详细原理请参考[原始论文](https://arxiv.org/abs/2309.00071)

    Args:
        config ([`~transformers.PretrainedConfig`]):
            模型配置对象，包含RoPE相关的所有参数设置
        device (`torch.device`):
            用于初始化逆频率张量的设备(CPU/GPU)
        seq_len (`int`, *optional*):
            当前序列长度，在YaRN中此参数未使用

    Returns:
        Tuple[torch.Tensor, float]: 返回一个元组，包含：
            - inv_freq: RoPE嵌入的逆频率张量，用于计算旋转位置编码
            - attention_factor: 应用于计算得到的cos/sin的后处理缩放因子，用于调整注意力权重
    """
    # 从配置中获取基础参数
    base = config.rope_config.rope_theta  # RoPE的基础频率参数，通常为10000
    partial_rotary_factor = config.rope_config.partial_rotary_factor  # 部分旋转因子，控制应用RoPE的维度比例
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)  # 注意力头的维度
    dim = int(head_dim * partial_rotary_factor)  # 实际应用RoPE的维度

    # YaRN特定的配置参数
    factor = config.rope_config.factor  # 序列长度扩展因子
    attention_factor = config.rope_config.attention_factor  # 注意力缩放因子
    mscale = config.rope_config.mscale  # 多尺度缩放参数
    mscale_all_dim = config.rope_config.mscale_all_dim  # 是否对所有维度应用多尺度缩放
    original_max_position_embeddings = config.original_max_position_embeddings  # 原始最大位置嵌入长度

    # 注意：DeepSeek-V3（以及其他可能的模型）会修改`max_position_embeddings`参数，
    # 并使用`original_max_position_embeddings`字段存储预训练时的原始值。
    # 它们使用这两个值的比值来计算默认的注意力缩放因子，而不是直接使用`factor`参数。
    if original_max_position_embeddings:
        # 如果存在原始最大位置嵌入长度，使用当前长度与原始长度的比值作为缩放因子
        factor = config.max_position_embeddings / original_max_position_embeddings
    else:
        # 否则使用当前的最大位置嵌入长度作为原始值
        original_max_position_embeddings = config.max_position_embeddings

    def get_mscale(scale, mscale=1.0):
        """
        计算多尺度缩放因子，用于调整注意力权重以适应长序列
        基于对数函数的缩放策略，当scale<=1时返回1.0，避免不必要的缩放

        Args:
            scale: 缩放比例，通常是序列长度扩展因子
            mscale: 多尺度缩放系数，默认为1.0

        Returns:
            float: 计算得到的缩放因子
        """
        if scale <= 1:
            return 1.0  # 当缩放比例<=1时，不进行缩放
        return 0.1 * mscale * math.log(scale) + 1.0  # 使用对数函数进行平滑缩放

    # 根据论文建议设置注意力缩放因子
    if attention_factor is None:
        if mscale and mscale_all_dim:
            # 如果启用了多尺度缩放且对所有维度应用，计算两个缩放因子的比值
            attention_factor = float(get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim))
        else:
            # 否则直接使用基础缩放因子
            attention_factor = get_mscale(factor)

    # 可选配置选项
    # beta_fast/beta_slow: 论文中建议的参数，默认值分别为32/1
    # 这些参数用于控制频率调整的范围和速度
    beta_fast = config.rope_config.beta_fast or 32  # 快速调整因子，控制高频部分的调整
    beta_slow = config.rope_config.beta_slow or 1  # 慢速调整因子，控制低频部分的调整

    # 计算逆频率的核心函数定义
    def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
        """
        根据旋转次数反向计算维度的公式
        用于确定在给定旋转次数下应该调整哪些维度

        Args:
            num_rotations: 旋转次数
            dim: 总维度
            base: 基础频率
            max_position_embeddings: 最大位置嵌入长度

        Returns:
            float: 计算得到的修正维度
        """
        return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))

    def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
        """
        基于旋转次数寻找维度范围的边界
        确定需要进行频率调整的维度区间

        Args:
            low_rot: 低频旋转次数
            high_rot: 高频旋转次数
            dim: 总维度
            base: 基础频率
            max_position_embeddings: 最大位置嵌入长度

        Returns:
            Tuple[int, int]: 维度范围的上下界
        """
        low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
        high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
        return max(low, 0), min(high, dim - 1)  # 确保范围在有效维度内

    def linear_ramp_factor(min, max, dim):
        """
        计算线性斜坡因子，用于在指定维度范围内实现平滑过渡
        通过线性插值在min和max之间创建一个0到1的斜坡函数

        Args:
            min: 最小维度
            max: 最大维度
            dim: 总维度数

        Returns:
            torch.Tensor: 线性斜坡因子张量，值域[0, 1]
        """
        if min == max:
            max += 0.001  # 防止除零错误，避免奇点

        # 创建线性函数并限制在[0, 1]范围内
        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
        ramp_func = torch.clamp(linear_func, 0, 1)
        return ramp_func

    # 变量命名说明：这里的"interpolation"（插值）来源于原始技术，
    # 其中我们通过插值位置ID来扩展可能的上下文长度。
    # 换句话说，插值 = 应用缩放因子。

    # 计算基础位置频率
    pos_freqs = base ** (torch.arange(0, dim, 2).to(device=device, dtype=torch.float) / dim)

    # 计算外推和插值的逆频率
    inv_freq_extrapolation = 1.0 / pos_freqs  # 外推逆频率（原始频率）
    inv_freq_interpolation = 1.0 / (factor * pos_freqs)  # 插值逆频率（缩放后频率）

    # 确定需要修正的维度范围
    low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings)

    # 获取经过外推修正的n维旋转缩放因子
    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).to(device=device, dtype=torch.float)

    # 通过加权组合插值和外推的逆频率，得到最终的逆频率
    # 这种组合方式能够在不同维度上应用不同程度的缩放，实现平滑的频率调整
    inv_freq = (
            inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)  # 插值部分权重
            + inv_freq_extrapolation * inv_freq_extrapolation_factor  # 外推部分权重
    )

    return inv_freq, attention_factor


ROPE_INIT_FUNCTIONS = {
    "default": _compute_default_rope_parameters,
    "dynamic": _compute_dynamic_ntk_parameters,
    "yarn": _compute_yarn_parameters,
}


def rotate_half(x):
    """
    对输入张量的后半部分隐藏维度进行旋转操作。

    这是RoPE (Rotary Position Embedding) 实现中的核心辅助函数，用于实现二维复数旋转。
    通过将张量分成两半并对后半部分取负，然后交换位置，实现了复数平面中的90度旋转效果。
    这种旋转操作是RoPE能够编码位置信息的关键机制之一。

    数学原理：
    对于复数 z = a + bi，其中 a 是实部，b 是虚部
    旋转90度后得到：z' = -b + ai
    这个函数实现了对应的张量操作

    Args:
        x (torch.Tensor): 输入张量，形状为 [..., dim]，其中 dim 是隐藏维度
                        通常是偶数，因为需要平均分成两部分

    Returns:
        torch.Tensor: 旋转后的张量，形状与输入相同
                     后半部分被取负并与前半部分交换位置

    Example:
        如果输入 x = [a, b, c, d] (dim=4)
        则 x1 = [a, b] (前半部分)
             x2 = [c, d] (后半部分)
        输出为 [-c, -d, a, b]
    """
    # 将输入张量沿着最后一个维度分成两半
    # x1: 前半部分维度 [..., :dim//2]
    x1 = x[..., : x.shape[-1] // 2]

    # x2: 后半部分维度 [..., dim//2:]
    x2 = x[..., x.shape[-1] // 2:]

    # 将后半部分取负，然后与前半部分拼接
    # 实现了 [x1, x2] -> [-x2, x1] 的变换
    # 这对应于复数平面中的90度逆时针旋转
    return torch.cat((-x2, x1), dim=-1)


def rotate_half(x):
    """
    对输入张量的后半部分隐藏维度进行旋转操作。

    这是RoPE (Rotary Position Embedding) 实现中的核心辅助函数，用于实现二维复数旋转。
    通过将张量分成两半并对后半部分取负，然后交换位置，实现了复数平面中的90度旋转效果。
    这种旋转操作是RoPE能够编码位置信息的关键机制之一。

    数学原理：
    对于复数 z = a + bi，其中 a 是实部，b 是虚部
    旋转90度后得到：z' = -b + ai
    这个函数实现了对应的张量操作

    Args:
        x (torch.Tensor): 输入张量，形状为 [..., dim]，其中 dim 是隐藏维度
                        通常是偶数，因为需要平均分成两部分

    Returns:
        torch.Tensor: 旋转后的张量，形状与输入相同
                     后半部分被取负并与前半部分交换位置

    Example:
        如果输入 x = [a, b, c, d] (dim=4)
        则 x1 = [a, b] (前半部分)
             x2 = [c, d] (后半部分)
        输出为 [-c, -d, a, b]
    """
    # 将输入张量沿着最后一个维度分成两半
    # x1: 前半部分维度 [..., :dim//2]
    x1 = x[..., : x.shape[-1] // 2]

    # x2: 后半部分维度 [..., dim//2:]
    x2 = x[..., x.shape[-1] // 2:]

    # 将后半部分取负，然后与前半部分拼接
    # 实现了 [x1, x2] -> [-x2, x1] 的变换
    # 这对应于复数平面中的90度逆时针旋转
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    """
    对查询(Query)和键(Key)张量应用旋转位置编码(Rotary Position Embedding)。

    这是RoPE实现的核心应用函数，通过复数旋转的方式将位置信息注入到注意力机制的查询和键向量中。
    RoPE的优势在于它能够相对位置编码，使得模型能够更好地处理序列中的位置关系，同时保持外推能力。

    数学原理：
    对于位置m，RoPE将向量q_m旋转为：R_Θ,m * q_m
    其中R_Θ,m是对角旋转矩阵，通过cos和sin实现复数旋转

    Args:
        q (torch.Tensor): 查询张量，形状通常为 [batch_size, num_heads, seq_len, head_dim]
                        或 [batch_size, seq_len, num_heads, head_dim]
        k (torch.Tensor): 键张量，形状与查询张量相同
        cos (torch.Tensor): 旋转编码的余弦部分，形状通常为 [seq_len, head_dim]
                         或 [batch_size, seq_len, head_dim]
        sin (torch.Tensor): 旋转编码的正弦部分，形状与cos相同
        position_ids (torch.Tensor, *optional*):
            已弃用且未使用的参数，保留是为了向后兼容
        unsqueeze_dim (int, *optional*, defaults to 1):
            指定在哪个维度上对cos和sin进行unsqueeze操作，以便能够正确地广播到q和k的维度。

            广播机制说明：
            - cos和sin的形状通常是 [batch_size, seq_len, head_dim]
            - 如果q和k的形状是 [batch_size, heads, seq_len, head_dim]，设置unsqueeze_dim=1
              使cos/sin变为 [batch_size, 1, seq_len, head_dim]，可以正确广播
            - 如果q和k的形状是 [batch_size, seq_len, heads, head_dim]，设置unsqueeze_dim=2
              使cos/sin变为 [batch_size, seq_len, 1, head_dim]，可以正确广播

    Returns:
        tuple(torch.Tensor): 包含两个元素的元组：
            - q_embed: 应用旋转位置编码后的查询张量，形状与输入q相同
            - k_embed: 应用旋转位置编码后的键张量，形状与输入k相同

    Example:
        q = torch.randn(2, 8, 128, 64)  # [batch, heads, seq_len, head_dim]
        k = torch.randn(2, 8, 128, 64)  # [batch, heads, seq_len, head_dim]
        cos = torch.randn(128, 64)      # [seq_len, head_dim]
        sin = torch.randn(128, 64)      # [seq_len, head_dim]

        q_rot, k_rot = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
        # q_rot, k_rot 形状仍为 [2, 8, 128, 64]
    """
    # 对cos和sin张量在指定维度上进行unsqueeze操作，增加一个维度以便广播
    # 这是实现张量形状匹配的关键步骤
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)

    # 对查询张量应用旋转位置编码
    # 使用复数旋转公式：q_rot = q * cos + rotate_half(q) * sin
    # 其中：
    # - q * cos: 实部乘法（保持原始方向）
    # - rotate_half(q) * sin: 虚部乘法（旋转90度后的分量）
    # 这个公式实现了二维复数平面中的旋转操作
    q_embed = (q * cos) + (rotate_half(q) * sin)

    # 对键张量应用相同的旋转位置编码
    # 使用与查询张量完全相同的旋转操作，确保位置编码的一致性
    k_embed = (k * cos) + (rotate_half(k) * sin)

    # 返回旋转后的查询和键张量
    # 这些张量现在包含了位置信息，可以用于后续的注意力计算
    return q_embed, k_embed


class RotaryEmbedding(nn.Module):
    """
    旋转位置编码(Rotary Position Embedding, RoPE)的PyTorch实现类。

    RoPE是一种先进的位置编码技术，通过复数旋转的方式将位置信息注入到注意力机制的查询和键向量中。
    相比传统的绝对位置编码，RoPE具有更好的外推能力和相对位置感知能力，特别适合处理长序列任务。

    该类支持多种RoPE变体：
    - default: 标准RoPE实现
    - dynamic: 动态NTK RoPE，支持推理时动态调整序列长度
    - yarn: YaRN RoPE，改进的长序列外推方法

    主要特性：
    1. 支持多种RoPE算法类型
    2. 动态频率更新（针对dynamic类型）
    3. 高效的张量计算和内存管理
    4. 自动处理设备兼容性和数值精度

    Attributes:
        rope_type (str): RoPE类型，支持"default"、"dynamic"、"yarn"
        max_seq_len_cached (int): 当前缓存的序列长度
        original_max_seq_len (int): 原始最大序列长度
        attention_scaling (float): 注意力缩放因子
        inv_freq (torch.Tensor): 逆频率张量，用于计算旋转编码
    """

    def __init__(
            self,
            config: Optional[Config] = None,
            device=None,
    ):
        """
        初始化RotaryEmbedding模块。

        Args:
            config (Optional[Config]): 模型配置对象，包含RoPE相关参数
            device: 初始化设备(CPU/GPU)，如果为None则使用默认设备
        """
        super().__init__()

        # 确定RoPE类型：从配置中获取，如果未指定则使用默认类型
        if config is not None and config.rope_config.rope_type is not None:
            self.rope_type = config.rope_config.rope_type
        else:
            self.rope_type = "default"

        # 设置序列长度相关参数
        self.max_seq_len_cached = config.max_position_embeddings  # 当前缓存的序列长度
        self.original_max_seq_len = config.max_position_embeddings  # 原始最大序列长度

        # 保存配置和初始化函数引用
        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]  # 根据类型选择对应的初始化函数

        # 初始化逆频率张量和注意力缩放因子
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        # 注册为缓冲区，不作为模型参数参与训练
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq  # 保存原始逆频率，用于动态更新时的重置

    def _dynamic_frequency_update(self, position_ids, device):
        """
        动态更新频率参数，专门用于dynamic类型的RoPE。

        动态RoPE层需要在以下情况下重新计算`inv_freq`：
        1. 序列长度超过缓存的长度（允许缩放以适应更长的序列）
        2. 当前序列长度回到原始尺度（避免短序列时的精度损失）

        这种动态更新机制使得模型能够在推理时处理不同长度的序列，
        同时保持数值精度和外推能力。

        Args:
            position_ids (torch.Tensor): 位置ID张量，用于确定当前序列长度
            device: 计算设备
        """
        # 计算当前序列长度（最大位置ID + 1）
        seq_len = torch.max(position_ids) + 1

        # 情况1：序列长度增长超过缓存长度
        if seq_len > self.max_seq_len_cached:  # growth
            # 使用新的序列长度重新计算逆频率和注意力缩放因子
            inv_freq, self.attention_scaling = self.rope_init_fn(
                self.config, device, seq_len=seq_len
            )
            # 更新缓存的逆频率和最大序列长度
            self.register_buffer("inv_freq", inv_freq, persistent=False)
            self.max_seq_len_cached = seq_len

        # 情况2：序列长度回到原始尺度，且当前缓存长度大于原始长度
        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
            # 重置为原始逆频率，避免短序列时的精度损失
            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
            self.max_seq_len_cached = self.original_max_seq_len

    @torch.no_grad()
    def forward(self, x, position_ids):
        """
        前向传播计算旋转位置编码的cos和sin分量。

        这是RoPE的核心计算函数，根据位置ID和逆频率计算旋转编码，
        为后续的查询和键向量旋转提供必要的三角函数值。

        Args:
            x (torch.Tensor): 输入张量，用于确定数据类型和设备
            position_ids (torch.Tensor): 位置ID张量，形状为 [batch_size, seq_len]

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: 包含两个元素的元组：
                - cos: 余弦分量张量，形状为 [batch_size, seq_len, head_dim]
                - sin: 正弦分量张量，形状与cos相同
        """
        # 如果是dynamic类型的RoPE，首先进行动态频率更新
        if "dynamic" == self.rope_type:
            self._dynamic_frequency_update(position_ids, device=x.device)

        # ========== RoPE核心计算块 ==========

        # 扩展逆频率张量以匹配批次维度
        # inv_freq形状: [head_dim//2] -> [batch_size, head_dim//2, 1]
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)

        # 扩展位置ID张量以匹配频率维度
        # position_ids形状: [batch_size, seq_len] -> [batch_size, 1, seq_len]
        position_ids_expanded = position_ids[:, None, :].float()

        # 强制使用float32精度以避免数值精度问题
        # 参考链接：https://github.com/huggingface/transformers/pull/29285
        device_type = x.device.type
        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"

        # 在禁用自动混合精度的情况下计算频率
        with torch.autocast(device_type=device_type, enabled=False):
            # 计算旋转频率：inv_freq @ position_ids
            # 结果形状: [batch_size, 1, seq_len] -> 转置为 [batch_size, seq_len, 1]
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)

            # 复制频率以匹配完整的head_dim（RoPE只对偶数维度操作）
            # freqs形状: [batch_size, seq_len, 1] -> [batch_size, seq_len, 2] (通过拼接)
            emb = torch.cat((freqs, freqs), dim=-1)

            # 计算余弦和正弦分量
            cos = emb.cos()
            sin = emb.sin()

        # ========== 高级RoPE后处理 ==========

        # 对于高级RoPE类型（如yarn），应用后处理缩放因子
        # 这等效于对注意力权重进行缩放，用于改进长序列性能
        cos = cos * self.attention_scaling
        sin = sin * self.attention_scaling

        # 返回与输入数据类型匹配的cos和sin张量
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


def _expand_mask(
        mask: torch.Tensor,
        dtype: torch.dtype,
        tgt_len: Optional[int] = None
):
    """
    将注意力掩码从 `[bsz, seq_len]` 扩展为 `[bsz, 1, tgt_seq_len, src_seq_len]` 格式。

    该函数是Transformer注意力机制中的核心工具函数，用于将输入的注意力掩码
    扩展为适合多头注意力计算的4D张量格式。支持处理2D和3D输入掩码，
    并将布尔型或二进制掩码转换为适合注意力计算的数值型掩码。

    Args:
        mask (torch.Tensor): 输入的注意力掩码张量
            - 2D形状: [batch_size, src_seq_len] - 标准的填充掩码
            - 3D形状: [batch_size, src_seq_len, ...] - 扩展的掩码格式
            - 数值范围: 通常为0/1或True/False，表示有效/无效位置
        dtype (torch.dtype): 输出张量的数据类型
            - 作用: 确保掩码与注意力分数张量的数据类型一致
            - 常见值: torch.float16, torch.float32, torch.bfloat16
        tgt_len (Optional[int]): 目标序列长度
            - 默认值: None，此时使用src_len作为目标长度
            - 作用: 指定扩展后掩码的目标序列维度大小
            - 应用场景: 在解码器-编码器注意力中，目标序列长度可能与源序列不同

    Returns:
        torch.Tensor: 扩展后的注意力掩码张量
            - 形状: [batch_size, 1, tgt_seq_len, src_seq_len]
            - 数值特征: 有效位置为0，无效位置为该数据类型的最小值
            - 广播机制: 第2维为1，便于在多头注意力中广播到所有注意力头

    Examples:
        >>> mask = torch.tensor([[1, 1, 1, 0]])  # [bsz=1, seq_len=4]
        >>> expanded = _expand_mask(mask, torch.float16)
        >>> expanded.shape
        torch.Size([1, 1, 4, 4])
        >>> expanded[0, 0, :, :]  # 最后一个token被mask
        tensor([[0., 0., 0., -65504.],
                [0., 0., 0., -65504.],
                [0., 0., 0., -65504.],
                [0., 0., 0., -65504.]], dtype=torch.float16)
    """
    # 根据输入掩码的维度进行不同的处理逻辑
    if len(mask.size()) == 3:
        # 处理3D输入掩码的情况 [batch_size, src_seq_len, extra_dim]
        bsz, src_len, _ = mask.size()
        # 如果未指定目标长度，则使用源序列长度
        tgt_len = tgt_len if tgt_len is not None else src_len
        # 扩展掩码维度: [bsz, src_len, _] -> [bsz, 1, tgt_len, src_len]
        # 使用expand进行高效广播，避免内存复制
        expanded_mask = mask[:, None, :, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
    else:
        # 处理2D输入掩码的情况 [batch_size, src_seq_len]
        bsz, src_len = mask.size()
        # 如果未指定目标长度，则使用源序列长度
        tgt_len = tgt_len if tgt_len is not None else src_len
        # 扩展掩码维度: [bsz, src_len] -> [bsz, 1, tgt_len, src_len]
        # 先插入两个维度None，然后扩展到目标形状
        expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    # 掩码值转换: [true, true, true, false] -> [0, 0, 0, 1]
    # 将布尔型掩码转换为数值型，1表示需要被mask的位置
    inverted_mask = 1.0 - expanded_mask

    # 将mask位置填充为极小值，使其在softmax计算中概率趋近于0
    # 示例输出: tensor([[[[     0.,      0.,      0., -65504.],
    #                   [     0.,      0.,      0., -65504.],
    #                   [     0.,      0.,      0., -65504.],
    #                   [     0.,      0.,      0., -65504.]]]])
    # 其中-65504是float16的最小值，不同数据类型会有不同的最小值
    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)


def _make_causal_mask(
        input_ids_shape: torch.Size,
        dtype: torch.dtype,
        device:
        torch.device,
        past_key_values_length: int = 0
):
    """
    生成用于自回归解码的因果注意力掩码。

    该函数创建一个下三角矩阵掩码，确保在自注意力计算中，每个位置只能关注到
    自己和之前的位置，不能看到未来的信息。这是实现Transformer自回归解码的
    关键机制，防止信息泄露并确保模型按顺序生成token。

    数学原理：
    对于序列长度为L的输入，生成的因果掩码M满足：
    M[i,j] = 0, 如果 j <= i (允许关注)
    M[i,j] = -∞, 如果 j > i (禁止关注未来)

    Args:
        input_ids_shape (torch.Size): 输入张量的形状
            - 格式: (batch_size, seq_length)
            - batch_size: 批次大小，用于最终掩码的广播
            - seq_length: 目标序列长度，必须大于1
            - 约束: seq_length > 1，否则无法形成有效的因果掩码
        dtype (torch.dtype): 掩码张量的数据类型
            - 作用: 确保掩码与注意力分数张量的数据类型一致
            - 常见值: torch.float16, torch.float32, torch.bfloat16
            - 影响: 决定了mask中极小值的具体数值
        device (torch.device): 张量所在的计算设备
            - 作用: 确保所有张量在同一设备上进行计算
            - 常见值: 'cpu', 'cuda:0', 'cuda:1'等
            - 重要性: 避免跨设备计算带来的性能开销
        past_key_values_length (int, optional): 过去的键值对长度
            - 默认值: 0，表示没有历史键值对
            - 作用: 在增量推理时，考虑之前已经计算的键值对
            - 应用场景: 用于支持Transformer的增量解码和缓存机制

    Returns:
        torch.Tensor: 扩展后的因果注意力掩码
            - 形状: [batch_size, 1, tgt_len, tgt_len + past_key_values_length]
            - 数值特征: 
                * 允许关注的位置: 0
                * 禁止关注的位置: 该数据类型的最小值(如-65504 for float16)
            - 广播机制: 第2维为1，便于在多头注意力中广播到所有注意力头

    Examples:
        >>> input_shape = torch.Size([2, 4])  # batch_size=2, seq_len=4
        >>> mask = _make_causal_mask(input_shape, torch.float16, 'cpu')
        >>> mask.shape
        torch.Size([2, 1, 4, 4])
        >>> mask[0, 0, :, :]  # 第一个样本的因果掩码
        tensor([[0., -65504., -65504., -65504.],
                [0.,      0., -65504., -65504.],
                [0.,      0.,      0., -65504.],
                [0.,      0.,      0.,      0.]], dtype=torch.float16)
    """
    # 解包输入形状，获取批次大小和目标序列长度
    bsz, tgt_len = input_ids_shape

    # 创建初始掩码矩阵，填充为极小值
    # 形状: [tgt_len, tgt_len]，所有元素初始化为数据类型的最小值
    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)

    # 创建序列索引张量，用于生成下三角掩码的条件判断
    # 生成: (0, 1, 2, 3, 4, ...)，长度等于目标序列长度
    mask_cond = torch.arange(mask.size(-1), device=device)

    # 核心逻辑：通过广播比较生成下三角掩码
    # mask_cond shape (5) 广播到 (5, 5):
    #  [[0, 1, 2, 3, 4],
    #   [0, 1, 2, 3, 4],
    #   ...
    #   [0, 1, 2, 3, 4]]
    #
    # (mask_cond + 1).view(mask.size(-1), 1) shape (5, 1) 广播到 (5, 5):
    # [[1, 1, 1, 1, 1],
    #  [2, 2, 2, 2, 2],
    #  ...
    #  [5, 5, 5, 5, 5]]
    #
    # mask_cond < (mask_cond + 1).view(mask.size(-1), 1) 比较结果:
    # [[True, False, False, False, False],   # 第0行只能关注位置0
    #  [True, True, False, False, False],   # 第1行可以关注位置0,1
    #  [True, True, True, False, False],    # 第2行可以关注位置0,1,2
    #  [True, True, True, True, False],     # 第3行可以关注位置0,1,2,3
    #  [True, True, True, True, True]]      # 第4行可以关注所有位置
    #
    # 最终掩码形状 (seq_len, seq_len):
    #  [[0, -inf, -inf, -inf, -inf],   # 对角线及以下为0，以上为-inf
    #   [0, 0, -inf, -inf, -inf],
    #   ...
    #   [0, 0, 0, 0, 0]]
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)

    # 确保掩码数据类型与输入参数一致
    mask = mask.to(dtype)

    # 处理增量推理场景：如果有历史键值对，需要在掩码左侧添加零矩阵
    # 这允许当前token关注所有历史键值对，实现增量解码
    if past_key_values_length > 0:
        # 创建左侧零矩阵 [tgt_len, past_key_values_length]，所有元素为0
        # 表示当前token可以关注所有历史位置
        left_zeros = torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device)
        # 水平拼接: [历史部分 | 当前部分]
        mask = torch.cat([left_zeros, mask], dim=-1)

    # 扩展掩码维度以适应多头注意力计算的格式要求
    # [tgt_len, total_len] -> [bsz, 1, tgt_len, total_len]
    # 其中 total_len = tgt_len + past_key_values_length
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def prepare_decoder_attention_mask(
        attention_mask,
        doc_boundary_mask,
        input_shape,
        past_key_values_length,
        dtype,
        device
):
    """
    准备解码器注意力掩码，整合因果掩码、填充掩码和文档边界掩码。

    该函数是Transformer解码器中注意力掩码的核心处理函数，负责将多种类型的掩码
    整合为一个统一的注意力掩码。支持自回归解码的因果约束、序列填充处理以及
    文档边界限制，确保模型在生成过程中遵循正确的注意力模式。

    掩码整合策略：
    1. 因果掩码：确保每个位置只能关注自己和之前的位置
    2. 填充掩码：忽略填充token的影响
    3. 文档边界掩码：限制注意力在文档边界内

    Args:
        attention_mask (torch.Tensor, optional): 基础注意力掩码
            - 形状: [batch_size, seq_length] 或 [batch_size, seq_length, seq_length]
            - 作用: 标识有效token和填充token的位置
            - 数值: 通常为1(True)表示有效，0(False)表示填充
            - 示例: torch.ones((batch_size, seq_length), dtype=torch.bool)
        doc_boundary_mask (torch.Tensor, optional): 文档边界掩码
            - 形状: [batch_size, 1, tgt_seq_len, src_seq_len]
            - 作用: 限制注意力在文档边界内，防止跨文档信息泄露
            - 应用场景: 多文档处理、长文本分段等场景
            - 数值: 有效位置为0，边界限制位置为极小值
        input_shape (torch.Size): 输入张量的形状
            - 格式: (batch_size, seq_length)
            - batch_size: 批次大小，用于掩码的维度扩展
            - seq_length: 目标序列长度，决定因果掩码的大小
            - 约束: seq_length > 1时才生成因果掩码
        past_key_values_length (int): 过去的键值对长度
            - 作用: 支持增量推理，考虑历史计算的键值对
            - 默认值: 0，表示没有历史缓存
            - 应用场景: Transformer的增量解码、文本生成等
        dtype (torch.dtype): 输出掩码的数据类型
            - 作用: 确保与注意力分数张量的数据类型一致
            - 常见值: torch.float16, torch.float32, torch.bfloat16
            - 影响: 决定掩码中极小值的具体数值
        device (torch.device): 计算设备
            - 作用: 确保所有张量在同一设备上进行计算
            - 常见值: 'cpu', 'cuda:0', 'cuda:1'等
            - 重要性: 避免跨设备计算的性能开销

    Returns:
        torch.Tensor: 整合后的注意力掩码
            - 形状: [batch_size, 1, tgt_seq_len, src_seq_len]
            - 数值特征:
                * 允许关注的位置: 0
                * 禁止关注的位置: 极小值(如-65504 for float16)或-inf
            - 整合逻辑: 多种掩码的叠加效果，满足所有约束条件

    Examples:
        >>> input_shape = torch.Size([2, 4])  # batch_size=2, seq_len=4
        >>> attention_mask = torch.ones((2, 4), dtype=torch.bool)  # 全部有效
        >>> mask = prepare_decoder_attention_mask(
        ...     attention_mask, None, input_shape, 0, torch.float16, 'cpu'
        ... )
        >>> mask.shape
        torch.Size([2, 1, 4, 4])
        >>> mask[0, 0, :, :]  # 因果掩码效果
        tensor([[0., -65504., -65504., -65504.],
                [0.,      0., -65504., -65504.],
                [0.,      0.,      0., -65504.],
                [0.,      0.,      0.,      0.]], dtype=torch.float16)
    """
    # 初始化组合注意力掩码为None，用于后续的条件性掩码整合
    combined_attention_mask = None

    # 步骤1: 创建因果掩码（仅在序列长度大于1时）
    # 因果掩码确保自回归解码的正确性：每个位置只能关注自己和之前的位置
    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
    if input_shape[-1] > 1:
        # 调用_make_causal_mask生成下三角因果掩码
        # 示例输出: tensor([[[[     0., -65504., -65504., -65504.],
        #                   [     0.,      0., -65504., -65504.],
        #                   [     0.,      0.,      0., -65504.],
        #                   [     0.,      0.,      0.,      0.]]]])
        combined_attention_mask = _make_causal_mask(
            input_shape,
            dtype,
            device=device,
            past_key_values_length=past_key_values_length,
        )

    # 步骤2: 处理填充掩码（如果提供）
    # 填充掩码用于忽略序列中的填充token，防止模型关注无效位置
    if attention_mask is not None:
        # 使用_expand_mask将2D填充掩码扩展为4D格式
        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        # -65504 表示填充位置，在softmax中概率趋近于0
        # 示例输出: tensor([[[[     0.,      0.,      0., -65504.],
        #                   [     0.,      0.,      0., -65504.],
        #                   [     0.,      0.,      0., -65504.],
        #                   [     0.,      0.,      0., -65504.]]]])
        expanded_attn_mask = _expand_mask(attention_mask, dtype, tgt_len=input_shape[-1]).to(
            device
        )

        # 将填充掩码与因果掩码进行整合
        # 整合策略：如果只有一种掩码则直接使用，如果有多种则相加
        # 相加后的效果：同时满足因果约束和填充约束
        # 示例输出: tensor([[[[     0., -65504., -65504.,    -inf],
        #                   [     0.,      0., -65504.,    -inf],
        #                   [     0.,      0.,      0.,    -inf],
        #                   [     0.,      0.,      0., -65504.]]]], dtype=torch.float16)
        combined_attention_mask = (
            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
        )

    # 步骤3: 处理文档边界掩码（如果提供）
    # 文档边界掩码用于限制注意力在文档范围内，防止跨文档信息泄露
    if doc_boundary_mask is not None:
        # 将文档边界掩码与现有掩码进行整合
        # 整合逻辑与填充掩码相同：单一掩码直接使用，多种掩码相加
        # 这样可以确保同时满足因果约束、填充约束和文档边界约束
        combined_attention_mask = (
            doc_boundary_mask if combined_attention_mask is None else doc_boundary_mask + combined_attention_mask
        )

    # 返回最终整合的注意力掩码
    # 该掩码同时满足所有约束条件，可直接用于多头注意力计算
    return combined_attention_mask

class KVCache:
    """
    键值缓存类，用于大语言模型的增量推理优化

    KVCache(Key-Value Cache)是Transformer模型增量推理中的核心优化组件，
    通过缓存历史计算的键值对，避免在生成每个新token时重复计算整个序列的注意力，
    从而显著提高推理效率，将计算复杂度从O(n²)降低到O(n)。

    核心功能：
    1. 多层缓存管理：为每个Transformer层维护独立的键值缓存
    2. 动态扩展：支持缓存容量的动态增长，适应不同长度的序列
    3. 高效更新：通过张量拼接实现新token的快速缓存更新
    4. 序列长度跟踪：实时记录每个层的缓存序列长度

    技术特点：
    - 内存优化：按需分配缓存空间，避免预分配浪费
    - 索引管理：通过layer_idx实现多层缓存的精确访问
    - 张量操作：利用PyTorch的高效张量拼接操作
    - 边界处理：完善的空缓存和越界检查机制

    应用场景：
    - 文本生成：在自回归生成过程中避免重复计算
    - 对话系统：支持长对话的上下文维护
    - 代码补全：实现实时代码生成的高效推理
    """

    def __init__(self):
        """
        初始化KV缓存

        创建空的键值缓存列表，采用动态分配策略，
        只有在需要时才为特定层分配缓存空间，优化内存使用。

        缓存结构设计：
        - key_cache: 存储所有层的键状态张量列表
        - value_cache: 存储所有层的值状态张量列表
        - 每个层的缓存张量形状：(batch, num_heads, seq_len, head_size)
        """
        # 键缓存列表，按层索引存储每个层的键状态张量
        # 形状说明：(layer_idx, batch, num_heads, seq_len, head_size)
        # 其中layer_idx通过列表索引实现，batch/num_heads/seq_len/head_size为张量维度
        self.key_cache = []

        # 值缓存列表，与键缓存结构对称，存储每个层的值状态张量
        # 形状与键缓存完全一致，便于同步管理和访问
        self.value_cache = []

    def get_seq_len(self, layer_idx: Optional[int] = 0):
        """
        获取指定层的缓存序列长度

        用于增量推理中位置编码和注意力掩码的计算，
        确保新生成token的位置编码与历史缓存保持连续性。

        Args:
            layer_idx (Optional[int]): 层索引，指定要查询的Transformer层
                                     默认为0，即第一层
                                     支持多层模型的独立缓存管理

        Returns:
            int: 指定层的缓存序列长度
                 如果缓存为空或层不存在，返回0
                 否则返回该层缓存张量的seq_len维度大小

        检查逻辑：
        1. 全局缓存为空：整个KVCache还未存储任何数据
        2. 层索引越界：请求的层索引超出了当前缓存范围
        3. 层缓存为空：特定层的缓存还未初始化
        以上情况均返回0，表示无有效缓存数据
        """
        # 检查缓存是否为空的三种情况：
        # 1. 整个key_cache列表为空，表示没有任何层数据
        # 2. 请求的layer_idx超出了当前缓存列表的范围
        # 3. 指定层的缓存为空列表，表示该层还未初始化
        empty_cache = (
                len(self.key_cache) == 0
                or len(self.key_cache) <= layer_idx
                or len(self.key_cache[layer_idx]) == 0
        )

        # 根据缓存状态返回序列长度
        # 如果缓存非空，返回张量的倒数第二维(seq_len)的大小
        # 如果缓存为空，返回0表示无有效序列长度
        return self.key_cache[layer_idx].shape[-2] if not empty_cache else 0

    def update(self, key_states, value_states, layer_idx):
        """
        更新指定层的键值缓存

        这是KVCache的核心方法，负责将新计算的键值状态添加到缓存中，
        支持初始化、扩展和追加三种操作模式，确保缓存的连续性和完整性。

        Args:
            key_states (torch.Tensor): 新的键状态张量，形状为(batch, num_heads, seq_len, head_size)
                                     包含当前时间步计算得到的键向量
            value_states (torch.Tensor): 新的值状态张量，形状与key_states相同
                                       包含当前时间步计算得到的值向量
            layer_idx (int): 层索引，指定要更新的Transformer层
                           用于多层模型的独立缓存管理

        Returns:
            tuple[torch.Tensor, torch.Tensor]: 更新后的键值状态元组
                                             (key_cache[layer_idx], value_cache[layer_idx])
                                             返回完整的缓存张量，便于后续计算使用

        更新策略：
        1. 层索引越界：先填充中间层的空缓存，再添加新层数据
        2. 层缓存为空：直接初始化该层的缓存
        3. 层缓存存在：通过张量拼接在seq_len维度上追加新数据

        性能优化：
        - 使用torch.cat进行高效张量拼接
        - 按需分配内存，避免预分配浪费
        - 保持张量连续性，优化后续计算性能
        """
        # 情况1：层索引超出当前缓存范围
        # 需要先为中间层创建空缓存，然后添加新层数据
        if len(self.key_cache) <= layer_idx:
            # 为中间缺失的层创建空缓存占位
            # 确保缓存列表的连续性，便于后续访问
            for _ in range(len(self.key_cache), layer_idx):
                self.key_cache.append([])
                self.value_cache.append([])

            # 添加新层的键值缓存
            # 直接赋值，因为这是该层的第一次初始化
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)

        # 情况2：指定层的缓存为空
        # 表示该层已存在但还未初始化，直接赋值初始化
        elif len(self.key_cache[layer_idx]) == 0:
            self.key_cache[layer_idx] = key_states
            self.value_cache[layer_idx] = value_states

        # 情况3：指定层已有缓存数据
        # 需要在序列长度维度上拼接新的键值状态
        else:
            # 在seq_len维度(dim=-2)上拼接新旧键状态
            # 保持batch、num_heads、head_size维度不变
            self.key_cache[layer_idx] = torch.cat((self.key_cache[layer_idx], key_states), dim=-2)
            # 同样方式拼接值状态，确保键值缓存的一致性
            self.value_cache[layer_idx] = torch.cat((self.value_cache[layer_idx], value_states), dim=-2)

        # 返回更新后的完整键值缓存
        # 便于调用方直接使用最新的缓存数据进行后续计算
        return self.key_cache[layer_idx], self.value_cache[layer_idx]


class Attention(nn.Module):
    """
    多头注意力机制实现，支持分组查询注意力(GQA)和键值归一化。

    该类实现了现代大语言模型中的核心注意力机制，包含以下关键特性：
    1. 分组查询注意力(GQA)：通过减少键值头的数量来降低计算复杂度
    2. 可选的键值归一化：提高注意力计算的稳定性
    3. 多种注意力实现：支持SDPA(缩放点积注意力)和传统实现
    4. 旋转位置编码(RoPE)：集成位置信息到注意力计算中
    5. KV缓存：支持增量推理，提高生成效率

    架构设计：
    - 输入: [batch_size, seq_len, hidden_size]
    - Q/K/V投影: 通过线性层将隐藏状态映射到查询、键、值空间
    - 头部分组: 支持多头注意力和分组查询注意力
    - 位置编码: 应用RoPE旋转位置嵌入
    - 注意力计算: 使用SDPA或传统方法计算注意力权重
    - 输出投影: 将多头注意力结果映射回隐藏空间

    数学原理：
    Attention(Q,K,V) = softmax(QK^T/√d_k)V
    其中Q、K、V分别表示查询、键、值矩阵，d_k为键的维度
    """

    def __init__(self, config: Config, layer_idx: int):
        """
        初始化注意力层，配置模型参数和层索引。

        Args:
            config (Config): 模型配置对象，包含所有注意力相关的超参数
            layer_idx (int): 当前层的索引，用于KV缓存和MoE层判断
        """
        super().__init__()
        # 验证注意力头数必须是键值头数的整数倍，确保分组查询注意力的正确性
        assert config.num_attention_heads % config.num_key_value_heads == 0

        # 自动检测或手动指定注意力实现方式
        if config.attention_implementation == 'auto':
            # 自动检测是否支持SDPA(缩放点积注意力)，这是PyTorch 2.0+的优化实现
            self.use_sdpa_attention = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        else:
            # 手动指定使用SDPA实现
            self.use_sdpa_attention = config.attention_implementation == 'sdpa'

        # 配置键值归一化选项，提高注意力计算的数值稳定性
        self.use_qk_norm = config.use_qk_norm
        # 保存当前层索引，用于KV缓存管理和MoE层判断
        self.layer_idx = layer_idx

        # 基础模型参数配置
        self.hidden_size = config.hidden_size  # 隐藏层维度
        self.num_heads = config.num_attention_heads  # 查询头数量
        self.head_size = self.hidden_size // self.num_heads  # 每个注意力头的维度
        self.num_key_value_heads = config.num_key_value_heads  # 键值头数量(用于GQA)
        # 计算每组键值头对应的查询头数量，用于分组查询注意力
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        # 计算注意力缩放因子，防止梯度消失
        self.scale = self.head_size ** -0.5

        # 定义Q/K/V/O投影层，将隐藏状态映射到注意力空间
        # 查询投影：[hidden_size] -> [num_heads * head_size]
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_size, bias=False)
        # 键投影：[hidden_size] -> [num_key_value_heads * head_size]
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_size, bias=False)
        # 值投影：[hidden_size] -> [num_key_value_heads * head_size]
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_size, bias=False)
        # 输出投影：[num_heads * head_size] -> [hidden_size]
        self.o_proj = nn.Linear(self.num_heads * self.head_size, self.hidden_size, bias=False)
        # 注意力dropout层，用于正则化
        self.dropout = nn.Dropout(p=config.attention_dropout)

        # 如果启用键值归一化，创建Q/K的RMSNorm层
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_size)
            self.k_norm = RMSNorm(self.head_size)

    def _sdpa_kernel(self, enable_flash: bool = True, enable_math: bool = True, enable_mem_efficient: bool = True):
        """
        配置SDPA(缩放点积注意力)的后端实现选项。

        该方法根据PyTorch版本选择合适的SDPA后端，支持Flash Attention、
        数学实现和内存高效实现等多种优化算法。

        Args:
            enable_flash (bool): 是否启用Flash Attention优化
            enable_math (bool): 是否启用标准数学实现
            enable_mem_efficient (bool): 是否启用内存高效实现

        Returns:
            Context manager: SDPA内核的上下文管理器
        """
        # 根据PyTorch版本选择不同的SDPA配置方式
        if version.parse(torch.__version__).release < version.parse("2.3").release:
            # PyTorch 2.3之前的版本使用旧的SDPA配置接口
            return torch.backends.cuda.sdp_kernel(
                enable_flash=enable_flash, enable_math=enable_math, enable_mem_efficient=enable_mem_efficient
            )

        # PyTorch 2.3+版本使用新的SDPA配置接口
        backends = []
        if enable_flash:
            backends += [torch.nn.attention.SDPBackend.FLASH_ATTENTION]
        if enable_math:
            backends += [torch.nn.attention.SDPBackend.MATH]
        if enable_mem_efficient:
            backends += [torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION]
        return torch.nn.attention.sdpa_kernel(backends)

    def forward(
            self,
            hidden_states: torch.Tensor,
            position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            past_key_values: Optional[KVCache] = None
    ) -> torch.Tensor:
        """
        执行前向传播，计算多头注意力输出。

        该方法实现了完整的注意力计算流程，包括Q/K/V投影、头部分组、
        位置编码应用、KV缓存更新和注意力计算等步骤。

        Args:
            hidden_states (torch.Tensor): 输入隐藏状态
                - 形状: [batch_size, seq_len, hidden_size]
                - 作用: 模型的中间表示，包含序列的语义信息
            position_embeddings (Optional[Tuple[torch.Tensor, torch.Tensor]]): 旋转位置编码
                - 格式: (cos, sin) 元组
                - cos: 余弦位置编码，形状与query/key相同
                - sin: 正弦位置编码，形状与query/key相同
                - 作用: 为注意力计算注入位置信息
            attention_mask (Optional[torch.Tensor]): 注意力掩码
                - 形状: [batch_size, 1, tgt_seq_len, src_seq_len]
                - 作用: 控制注意力的可见性(因果掩码、填充掩码等)
            past_key_values (Optional[KVCache]): KV缓存
                - 作用: 存储历史键值对，支持增量推理
                - 应用场景: 文本生成、增量解码等

        Returns:
            torch.Tensor: 注意力层的输出
                - 形状: [batch_size, seq_len, hidden_size]
                - 作用: 包含注意力加权后的上下文信息
        """
        # 解包输入张量的形状信息
        batch, seq_len, _ = hidden_states.shape

        # 步骤1: Q/K/V投影 - 将隐藏状态映射到查询、键、值空间
        # 查询投影: [batch, seq_len, hidden_size] -> [batch, seq_len, num_heads*head_size]
        query_states = self.q_proj(hidden_states)
        # 键投影: [batch, seq_len, hidden_size] -> [batch, seq_len, num_key_value_heads*head_size]
        key_states = self.k_proj(hidden_states)
        # 值投影: [batch, seq_len, hidden_size] -> [batch, seq_len, num_key_value_heads*head_size]
        value_states = self.v_proj(hidden_states)

        # 步骤2: 头部分组 - 将投影后的状态重塑为多头格式
        # 查询重塑: [batch, seq_len, num_heads*head_size] -> [batch, seq_len, num_heads, head_size]
        # 键重塑: [batch, seq_len, num_key_value_heads*head_size] -> [batch, seq_len, num_key_value_heads, head_size]
        # 值重塑: [batch, seq_len, num_key_value_heads*head_size] -> [batch, seq_len, num_key_value_heads, head_size]
        query_states = query_states.reshape(batch, seq_len, -1, self.head_size)
        key_states = key_states.reshape(batch, seq_len, -1, self.head_size)
        value_states = value_states.reshape(batch, seq_len, -1, self.head_size)

        # 步骤3: 键值归一化(可选) - 提高注意力计算的数值稳定性
        if self.use_qk_norm:
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)

        # 步骤4: 维度重排 - 为后续的旋转位置编码和注意力计算准备维度顺序
        # 查询重排: [batch, seq_len, num_heads, head_size] -> [batch, num_heads, seq_len, head_size]
        # 键重排: [batch, seq_len, num_key_value_heads, head_size] -> [batch, num_key_value_heads, seq_len, head_size]
        # 值重排: [batch, seq_len, num_key_value_heads, head_size] -> [batch, num_key_value_heads, seq_len, head_size]
        query_states = query_states.permute(0, 2, 1, 3)
        key_states = key_states.permute(0, 2, 1, 3)
        value_states = value_states.permute(0, 2, 1, 3)

        # 步骤5: 应用旋转位置编码(RoPE) - 为注意力计算注入位置信息
        cos, sin = position_embeddings
        # 应用RoPE后，query和key包含了相对位置信息
        # query_states: [batch, num_heads, seq_len, head_size]
        # key_states: [batch, num_key_value_heads, seq_len, head_size]
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        # 步骤6: KV缓存更新(可选) - 支持增量推理
        if past_key_values is not None:
            # 将当前键值对更新到缓存中，用于后续的增量计算
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

        # 步骤7: 分组查询注意力(GQA)处理 - 处理键值头数少于查询头数的情况
        if self.num_key_value_heads == 1:
            # 单键值头情况(MQA)：键值头数为1，无需扩展
            # key_states, value_states: [batch, 1, seq_len, head_size]
            pass
        else:
            # 多键值头情况(GQA)：需要将键值扩展到与查询头数匹配
            # 首先为键值添加分组维度: [batch, num_key_value_heads, seq_len, head_size] -> [batch, num_key_value_heads, 1, seq_len, head_size]
            key_states = key_states[:, :, None, :, :]
            value_states = value_states[:, :, None, :, :]

            # 扩展键值到分组数量: [batch, num_key_value_heads, num_key_value_groups, seq_len, head_size]
            # 其中num_key_value_groups = num_heads // num_key_value_heads
            key_states = key_states.expand(
                batch, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_size
            )
            value_states = value_states.expand(
                batch, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_size
            )

            # 重塑为完整的注意力头格式: [batch, num_heads, seq_len, head_size]
            key_states = key_states.reshape(
                batch, self.num_key_value_heads * self.num_key_value_groups, key_states.shape[-2], self.head_size
            )
            value_states = value_states.reshape(
                batch, self.num_key_value_heads * self.num_key_value_groups, value_states.shape[-2], self.head_size
            )

        # 步骤8: 注意力计算 - 使用SDPA或传统方法
        if self.use_sdpa_attention:
            # 使用优化的SDPA实现(缩放点积注意力)
            with self._sdpa_kernel():
                # 训练时应用dropout，推理时关闭
                dropout_p = self.dropout.p if self.training else 0.0
                # 调用PyTorch优化的scaled_dot_product_attention函数
                attn = F.scaled_dot_product_attention(
                    query=query_states,
                    key=key_states,
                    value=value_states,
                    dropout_p=dropout_p,
                    attn_mask=attention_mask
                )
            # SDPA输出维度为[batch, num_heads, seq_len, head_size]，需要转置
            attn = attn.transpose(1, 2)
        else:
            # 使用传统注意力实现(手动计算)
            # 计算注意力分数: [batch, num_heads, seq_len, head_size] @ [batch, num_heads, head_size, seq_len]
            # 结果形状: [batch, num_heads, q_seq_len, k_seq_len]
            attn_scores = (self.scale * query_states) @ key_states.transpose(-1, -2)

            # 应用注意力掩码(因果掩码、填充掩码等)
            if attention_mask is not None:
                attn_scores = attn_scores + attention_mask

            # 计算注意力权重(softmax归一化)
            attn_weights = attn_scores.softmax(-1)

            # 训练时应用dropout正则化
            if self.training:
                attn_weights = self.dropout(attn_weights)

            # 计算注意力输出: [batch, num_heads, seq_len, seq_len] @ [batch, num_heads, seq_len, head_size]
            # 结果形状: [batch, num_heads, seq_len, head_size]
            attn = attn_weights @ value_states

            # 重排维度为: [batch, seq_len, num_heads, head_size]
            attn = attn.permute(0, 2, 1, 3)

        # 步骤9: 多头合并 - 将所有注意力头的输出合并
        # 重塑为: [batch, seq_len, num_heads*head_size]
        attn = attn.reshape(batch, seq_len, -1)

        # 步骤10: 输出投影 - 将注意力结果映射回隐藏空间
        # [batch, seq_len, num_heads*head_size] -> [batch, seq_len, hidden_size]
        out = self.o_proj(attn)

        return out


class MoEGate(nn.Module):
    """
    混合专家门控网络，实现MoE架构中的专家选择和路由机制

    MoEGate(Mixture of Experts Gate)是稀疏混合专家模型的核心组件，
    负责根据输入特征动态选择最合适的专家网络进行处理。
    通过可学习的门控机制，实现计算资源的高效分配和模型容量的扩展。

    核心功能：
    1. 专家选择：基于输入特征计算每个专家的权重分数
    2. Top-K路由：选择权重最高的K个专家进行计算
    3. 权重归一化：确保所选专家的权重和为1，保持数值稳定性
    4. 负载均衡：通过辅助损失实现专家间的负载均衡

    技术特点：
    - 动态路由：根据输入内容自适应选择专家
    - 稀疏激活：每个token只由少数专家处理，降低计算成本
    - 可微分：门控网络完全可微，支持端到端训练
    - 负载均衡：防止某些专家过载而其他专家闲置

    数学原理：
    门控网络使用线性变换 + softmax计算专家权重：
    g(x) = softmax(W_g * x)
    其中W_g是可学习的门控权重矩阵

    应用场景：
    - 大规模语言模型：扩展模型容量而不线性增加计算量
    - 多任务学习：不同专家处理不同类型的任务
    - 条件计算：根据输入动态调整计算路径
    """

    def __init__(self, config: Config):
        """
        初始化MoE门控网络

        根据配置参数设置门控网络的结构和超参数，
        包括专家数量、Top-K选择、权重归一化等关键配置。

        Args:
            config (Config): 模型配置对象，包含MoE相关的所有超参数
                           - moe_config.num_experts_per_tok: 每个token选择的专家数量
                           - moe_config.n_routed_experts: 路由专家总数
                           - moe_config.routed_scaling_factor: 路由缩放因子
                           - moe_config.seq_aux: 是否使用序列级辅助损失
                           - moe_config.norm_topk_prob: 是否归一化Top-K概率
                           - hidden_size: 输入特征的隐藏维度
        """
        super().__init__()
        self.config = config

        # 每个token选择的专家数量(Top-K)
        # 控制稀疏程度，K值越小计算越高效，但可能损失表达能力
        self.top_k = config.moe_config.num_experts_per_tok

        # 路由专家的总数量
        # 决定了MoE层的总容量，更多的专家提供更强的表达能力
        self.n_routed_experts = config.moe_config.n_routed_experts

        # 路由缩放因子
        # 用于调整门控权重的数值范围，影响专家选择的强度
        self.routed_scaling_factor = config.moe_config.routed_scaling_factor

        # 序列级辅助损失标志
        # 控制使用序列级还是token级的负载均衡损失计算
        self.seq_aux = config.moe_config.seq_aux

        # Top-K概率归一化标志
        # 当选择多个专家时，是否将其权重归一化为概率分布
        self.norm_topk_prob = config.moe_config.norm_topk_prob

        # 门控网络的输入维度
        # 与模型隐藏层维度一致，确保特征匹配
        self.gating_dim = config.hidden_size

        # 门控权重矩阵，核心可学习参数
        # 形状：(n_routed_experts, gating_dim)
        # 每行对应一个专家的权重向量，用于计算专家选择分数
        self.weight = nn.Parameter(
            torch.empty((self.n_routed_experts, self.gating_dim))
        )

        # 初始化门控权重
        # 使用Kaiming均匀初始化，适合ReLU类激活函数的前层
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """
        重置门控网络参数

        使用Kaiming均匀初始化方法初始化门控权重矩阵，
        这种初始化方式适合深度网络中的线性层，有助于梯度流动。

        数学原理：
        Kaiming初始化考虑了输入输出的维度，保持方差在传播过程中的稳定性
        W ~ U(-bound, bound), 其中bound = sqrt(6 / fan_in) * a
        这里a = sqrt(5)，是针对LeakyReLU的优化参数
        """
        import torch.nn.init as init
        # Kaiming均匀初始化，a=sqrt(5)是针对LeakyReLU的推荐值
        # 这种初始化有助于保持梯度的稳定性，避免梯度消失或爆炸
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, hidden_states):
        """
        MoE门控网络前向传播

        实现完整的专家选择流程：特征变换->专家评分->Top-K选择->权重归一化->负载均衡损失计算

        Args:
            hidden_states (torch.Tensor): 输入隐藏状态张量
                                       形状：(batch_size, seq_len, hidden_size)
                                       包含待处理的token特征表示

        Returns:
            tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
                - topk_idx: 选择的专家索引，形状：(batch_size*seq_len, top_k)
                          每个token选择的top_k个专家的索引
                - topk_weight: 专家权重，形状与topk_idx相同
                             每个选中专家的归一化权重
                - aux_loss: 辅助损失，仅在训练时返回
                          用于专家负载均衡的损失值，推理时为None

        处理流程：
        1. 特征展平：将3D输入转换为2D便于批量处理
        2. 门控计算：通过线性变换计算专家分数
        3. 概率转换：使用softmax将分数转换为概率分布
        4. Top-K选择：选择概率最高的K个专家
        5. 权重归一化：确保所选专家权重和为1
        6. 负载均衡：计算辅助损失促进专家均衡使用
        """
        # 获取输入张量的形状信息
        # bsz: 批次大小, seq_len: 序列长度, h: 隐藏维度
        bsz, seq_len, h = hidden_states.shape

        # 展平输入张量：(batch_size, seq_len, hidden_size) -> (batch_size*seq_len, hidden_size)
        # 便于统一处理所有token的门控计算
        hidden_states = hidden_states.view(-1, h)

        # 门控分数计算：线性变换 + 类型转换
        # 使用float32精度确保数值稳定性，避免梯度计算中的精度问题
        # 输出形状：(batch_size*seq_len, n_routed_experts)
        logits = F.linear(
            hidden_states.type(torch.float32), self.weight.type(torch.float32), None
        )

        # 将logits转换为概率分布
        # 在专家维度上应用softmax，确保所有专家权重和为1
        # 使用float32精度保证数值计算的准确性
        scores = logits.softmax(dim=-1, dtype=torch.float32)

        ### Top-K专家选择
        # 从所有专家中选择权重最高的K个专家
        # sorted=False表示不需要按权重排序，提高计算效率
        # topk_weight: 选中专家的原始权重值
        # topk_idx: 选中专家的索引位置
        topk_weight, topk_idx = torch.topk(
            scores, k=self.top_k, dim=-1, sorted=False
        )

        ### 权重归一化处理
        # 当选择多个专家时，需要将其权重归一化为概率分布
        if self.top_k > 1 and self.norm_topk_prob:
            # 计算分母：Top-K权重和 + 小常数(防止除零)
            # keepdim=True保持维度便于广播操作
            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
            # 归一化权重：使每个token的选中专家权重和为1
            # 乘以缩放因子调整权重范围
            topk_weight = topk_weight / denominator * self.routed_scaling_factor
        else:
            # 单专家或不归一化情况：直接应用缩放因子
            topk_weight = topk_weight * self.routed_scaling_factor

        ### 专家负载均衡辅助损失计算
        # 仅在训练阶段计算辅助损失，推理阶段返回None
        if self.training:
            scores_for_aux = scores
            aux_topk = self.top_k
            # 基于贪心Top-K方法计算辅助损失
            # 将专家索引重塑为(batch_size, seq_len * top_k)便于后续处理
            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)

            if self.seq_aux:
                ### 序列级辅助损失计算
                # 考虑序列内的专家使用分布，更细粒度的负载均衡
                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
                # 初始化专家计数矩阵：(batch_size, n_routed_experts)
                ce = torch.zeros(
                    bsz, self.n_routed_experts, device=hidden_states.device
                )
                # 使用scatter_add_统计每个专家被选择的次数
                # 在专家维度上累加选择计数
                ce.scatter_add_(
                    1,  # 在专家维度(dim=1)上进行累加
                    topk_idx_for_aux_loss,  # 专家索引
                    torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device),  # 计数权重
                ).div_(seq_len * aux_topk / self.n_routed_experts)  # 归一化计数

                # 计算辅助损失：专家使用频率 * 平均选择概率
                # 鼓励专家的均匀使用和合理的概率分配
                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean()
            else:
                ### Token级辅助损失计算
                # 简化的负载均衡损失，基于整体的专家使用统计
                # 将专家索引转换为one-hot编码矩阵
                mask_ce = F.one_hot(
                    topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts
                )
                # 计算专家选择频率：每个专家被选择的平均比例
                ce = mask_ce.float().mean(0)
                # 计算专家平均概率：所有token对每个专家的平均选择概率
                Pi = scores_for_aux.mean(0)
                # 计算专家负载：选择频率 * 专家数量
                fi = ce * self.n_routed_experts
                # 辅助损失：平均概率 * 负载，鼓励均衡使用
                aux_loss = (Pi * fi).sum()
        else:
            # 推理阶段不需要计算辅助损失
            aux_loss = None

        # 返回专家选择结果和辅助损失
        return topk_idx, topk_weight, aux_loss


class MoE(nn.Module):
    """
    混合专家模块，包含路由专家和共享专家的完整MoE实现

    MoE(Mixture of Experts)是一种稀疏条件计算架构，通过动态选择专家网络来处理输入，
    在保持模型表达能力的同时显著降低计算成本。该实现结合了路由专家和共享专家，
    提供了灵活且高效的混合专家解决方案。

    核心架构：
    1. 路由专家(Routed Experts)：通过门控网络动态选择的专家网络
    2. 共享专家(Shared Experts)：所有输入都会经过的共享网络
    3. 门控网络(Gate)：负责专家选择和权重分配

    技术特点：
    - 稀疏激活：每个token只由少数专家处理，降低计算复杂度
    - 动态路由：根据输入内容自适应选择最适合的专家
    - 负载均衡：通过辅助损失确保专家间的均衡使用
    - 双重结构：路由专家提供特异性，共享专家保证基础能力
    - 训练/推理优化：不同的计算策略适应不同场景

    数学原理：
    MoE的输出可以表示为：
    y = Σ(g_i(x) * Expert_i(x)) + Shared_Expert(x)
    其中g_i(x)是门控网络为第i个专家分配的权重

    应用场景：
    - 大规模语言模型：扩展模型容量而不线性增加计算量
    - 多模态学习：不同专家处理不同类型的数据
    - 个性化服务：根据用户特征动态选择处理路径
    """

    def __init__(self, config: Config, layer: Callable):
        """
        初始化混合专家模块

        根据配置参数创建路由专家、共享专家和门控网络，
        构建完整的MoE计算架构。

        Args:
            config (Config): 模型配置对象，包含MoE相关的所有超参数
                           - moe_config.num_experts_per_tok: 每个token选择的专家数量
                           - moe_config.n_routed_experts: 路由专家总数
                           - moe_config.intermediate_size: 专家网络的中间层维度
                           - moe_config.n_shared_experts: 共享专家数量，None表示不使用
            layer (Callable): 专家网络的构造函数
                             通常是MLP或其他前馈网络的类构造器
                             用于创建路由专家和共享专家实例
        """
        super().__init__()
        self.config = config

        # 每个token选择的专家数量(Top-K)
        # 控制稀疏程度，影响计算效率和模型表达能力
        self.num_experts_per_tok = config.moe_config.num_experts_per_tok

        # 每个rank(设备)上的路由专家数量
        # 在分布式训练中，决定每个设备负责的专家数量
        self.experts_per_rank = config.moe_config.n_routed_experts

        # 创建路由专家列表
        # 每个专家都是独立的神经网络，处理特定类型的输入特征
        # 使用ModuleList确保专家网络的正确注册和管理
        self.experts = nn.ModuleList(
            layer(config, intermediate_size=config.moe_config.intermediate_size)
            for _ in range(config.moe_config.n_routed_experts)
        )

        # 创建门控网络
        # 负责根据输入特征动态选择专家并分配权重
        self.gate = MoEGate(config)

        # 条件性创建共享专家
        # 共享专家为所有输入提供基础处理能力
        # 与路由专家形成互补，确保模型的基础表达能力
        if config.moe_config.n_shared_experts is not None:
            # 共享专家的中间层维度：单个专家维度 * 共享专家数量
            # 通过扩大维度来匹配多个专家的总容量
            intermediate_size = config.moe_config.intermediate_size * config.moe_config.n_shared_experts
            self.shared_experts = layer(config, intermediate_size=intermediate_size)

    def forward(self, hidden_states):
        """
        MoE模块的前向传播

        实现完整的混合专家计算流程：门控计算->专家选择->专家计算->结果聚合->共享专家融合

        Args:
            hidden_states (torch.Tensor): 输入隐藏状态张量
                                       形状：(batch_size, seq_len, hidden_size)
                                       包含待处理的token特征表示

        Returns:
            tuple[torch.Tensor, Optional[torch.Tensor]]:
                - y: MoE模块的输出张量，形状与输入hidden_states相同
                  包含路由专家和共享专家的融合结果
                - aux_loss: 辅助损失，用于专家负载均衡
                          仅在训练时返回有效值，推理时为None

        处理流程：
        1. 保存原始状态和形状信息
        2. 通过门控网络获取专家选择和权重
        3. 根据训练/推理模式选择不同的计算策略
        4. 路由专家结果聚合
        5. 共享专家计算和融合
        6. 返回最终结果和辅助损失
        """
        # 保存输入的原始状态，用于后续残差连接
        identity = hidden_states

        # 保存原始形状，用于最终输出重塑
        orig_shape = hidden_states.shape

        # 通过门控网络获取专家选择信息
        # topk_idx: 选择的专家索引
        # topk_weight: 专家权重
        # aux_loss: 负载均衡辅助损失
        topk_idx, topk_weight, aux_loss = self.gate(hidden_states)

        # 展平输入张量：(batch_size, seq_len, hidden_size) -> (batch_size*seq_len, hidden_size)
        # 便于统一处理所有token的专家分配
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

        # 展平专家索引：(batch_size, seq_len, top_k) -> (batch_size*seq_len*top_k)
        # 用于后续的专家分配和计算
        flat_topk_idx = topk_idx.view(-1)

        # 根据训练/推理模式选择不同的计算策略
        if self.training:
            ### 训练模式：使用动态路由和专家并行计算
            # 重复输入token：每个token复制top_k次，对应选择的专家数量
            # 形状：(batch_size*seq_len, hidden_size) -> (batch_size*seq_len*top_k, hidden_size)
            hidden_states = hidden_states.repeat_interleave(
                self.num_experts_per_tok, dim=0
            )

            # 初始化输出张量，用于存储专家计算结果
            y = torch.empty_like(hidden_states)

            # 遍历所有专家，处理分配给该专家的token
            for i, expert in enumerate(self.experts):
                # 找出分配给当前专家的token
                expert_mask = (flat_topk_idx == i)
                if expert_mask.any():  # 如果有token分配给该专家
                    # 使用当前专家处理对应的token
                    y[expert_mask] = expert(hidden_states[expert_mask])

            # 专家结果聚合：加权求和
            # 1. 重塑输出：(batch_size*seq_len*top_k, hidden_size) -> (batch_size*seq_len, top_k, hidden_size)
            # 2. 应用专家权重：(batch_size*seq_len, top_k, hidden_size) * (batch_size*seq_len, top_k, 1)
            # 3. 在专家维度求和：(batch_size*seq_len, hidden_size)
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)

            # 恢复数据类型和原始形状
            y = y.to(hidden_states.dtype).view(*orig_shape)
        else:
            ### 推理模式：使用优化的推理计算路径
            # moe_infer方法针对推理场景进行了优化，减少内存使用和提高计算效率
            y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)

        # 共享专家处理：如果配置了共享专家
        # 共享专家为所有输入提供基础处理能力，与路由专家形成互补
        if self.config.moe_config.n_shared_experts is not None:
            # 将路由专家结果与共享专家结果相加
            # 残差连接方式，增强模型的表达能力
            y = y + self.shared_experts(identity)

        # 返回最终结果和辅助损失
        return y, aux_loss

    @torch.no_grad()
    def moe_infer(self, x, topk_ids, topk_weight):
        """
        优化的MoE推理计算方法

        针对推理场景优化的专家计算实现，通过token分组和批量处理提高计算效率，
        减少内存碎片和计算开销。使用no_grad装饰器避免梯度计算，提高推理速度。

        Args:
            x (torch.Tensor): 展平后的输入张量，形状：(batch_size*seq_len, hidden_size)
            topk_ids (torch.Tensor): 专家索引张量，形状：(batch_size*seq_len, top_k)
            topk_weight (torch.Tensor): 专家权重张量，形状与topk_ids相同

        Returns:
            torch.Tensor: MoE推理结果，形状：(batch_size*seq_len, hidden_size)

        优化策略：
        1. Token分组：将分配给同一专家的token聚集在一起
        2. 批量处理：每个专家一次性处理所有分配的token
        3. 内存优化：减少中间张量的创建和复制
        4. 计算复用：避免重复计算相同的专家处理

        处理流程：
        1. 统计每个专家的token数量
        2. 按专家对token进行排序和分组
        3. 批量处理每个专家的token
        4. 重新排序并加权聚合结果
        """
        # 统计每个专家被分配的token数量
        # 创建计数矩阵：(batch_size*seq_len, n_experts)
        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
        # 使用scatter_在专家维度上累加计数
        cnts.scatter_(1, topk_ids, 1)
        # 计算每个专家的总token数量
        tokens_per_expert = cnts.sum(dim=0)

        # 对token按专家索引进行排序
        # 将分配给同一专家的token聚集在一起，便于批量处理
        idxs = topk_ids.view(-1).argsort()
        # 根据排序索引重新排列输入token
        # 注意：这里使用整数除法来处理top_k维度
        sorted_tokens = x[idxs // topk_ids.shape[1]]

        # 将token数量转换为numpy数组便于迭代
        tokens_per_expert = tokens_per_expert.cpu().numpy()

        # 批量处理每个专家的token
        outputs = []
        start_idx = 0
        for i, num_tokens in enumerate(tokens_per_expert):
            end_idx = start_idx + num_tokens
            if num_tokens == 0:
                # 如果没有token分配给该专家，跳过处理
                continue

            # 获取当前专家
            expert = self.experts[i]
            # 提取分配给该专家的token
            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
            # 专家批量处理token
            expert_out = expert(tokens_for_this_expert)
            # 收集专家输出
            outputs.append(expert_out)
            # 更新起始索引
            start_idx = end_idx

        # 拼接所有专家的输出结果
        # 如果有专家输出则拼接，否则返回空张量
        outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)

        # 将专家输出重新排序到原始token顺序
        # 创建输出张量并按原始索引重新排列
        new_x = torch.empty_like(outs)
        new_x[idxs] = outs

        # 加权聚合专家结果
        # 1. 重塑为专家维度：(batch_size*seq_len, top_k, hidden_size)
        # 2. 转换为权重数据类型并应用权重
        # 3. 在专家维度求和
        # 4. 恢复原始数据类型
        final_out = (
            new_x.view(*topk_ids.shape, -1)
            .type(topk_weight.dtype)
            .mul_(topk_weight.unsqueeze(dim=-1))
            .sum(dim=1)
            .type(new_x.dtype)
        )

        return final_out


class DecoderLayer(nn.Module):
    """
    Transformer解码器层实现，支持密集层和混合专家(MoE)架构。

    该类实现了现代大语言模型中的标准解码器层结构，包含以下关键组件：
    1. 多头自注意力机制：捕获序列内部的依赖关系
    2. 前馈神经网络：提供非线性变换能力
    3. 层归一化：稳定训练过程，加速收敛
    4. 残差连接：缓解梯度消失问题，支持深层网络训练
    5. 混合专家(MoE)：条件性支持稀疏专家网络架构

    架构设计：
    - 输入: [batch_size, seq_len, hidden_size]
    - 注意力子层: 自注意力 + 残差连接 + 层归一化
    - MLP子层: 前馈网络 + 残差连接 + 层归一化
    - 输出: [batch_size, seq_len, hidden_size] + 可选的辅助损失

    MoE支持：
    - 根据配置和层索引动态选择使用密集MLP或MoE
    - 支持路由专家和共享专家的混合架构
    - 在深层网络中启用MoE以提高模型容量和效率

    数学原理：
    DecoderLayer(x) = x + MLP(LayerNorm(x + Attention(LayerNorm(x))))
    其中Attention表示多头自注意力，MLP表示前馈神经网络
    """

    def __init__(self, config: Config, layer_idx: int):
        """
        初始化解码器层，配置注意力机制和前馈网络。

        根据配置参数和层索引，动态决定使用密集MLP还是混合专家(MoE)架构。
        MoE通常在网络的深层使用，以在保持计算效率的同时增加模型容量。

        Args:
            config (Config): 模型配置对象，包含所有解码器层相关的超参数
                - hidden_size: 隐藏层维度
                - moe_config: 混合专家配置，包含专家数量、路由参数等
                - n_dense_layer: 密集层数量，在此之前的层使用密集MLP
            layer_idx (int): 当前层的索引
                - 作用: 用于判断是否启用MoE架构
                - 范围: 从0开始，通常深层网络(如layer_idx >= n_dense_layer)使用MoE
        """
        super().__init__()

        # 初始化注意力机制的层归一化
        # RMSNorm是一种改进的层归一化，计算更高效，训练更稳定
        self.attn_norm = RMSNorm(config.hidden_size)

        # 初始化多头注意力机制
        # 传入配置和层索引，支持KV缓存和分组查询注意力
        self.attn = Attention(config, layer_idx)

        # 初始化MLP的层归一化
        # 为前馈网络提供归一化，稳定训练过程
        self.mlp_norm = RMSNorm(config.hidden_size)

        # 判断是否使用混合专家(MoE)架构
        # 通过多个条件检查确保MoE配置的完整性和适用性
        use_moe = (
            # 检查MoE配置是否存在
                config.moe_config
                # 检查中间层维度是否配置(专家网络的隐藏层大小)
                and config.moe_config.intermediate_size
                # 检查每token选择的专家数量是否配置
                and config.moe_config.num_experts_per_tok
                # 检查路由专家数量是否配置(可路由的专家总数)
                and config.moe_config.n_routed_experts
                # 检查共享专家数量是否配置(始终激活的专家)
                and config.moe_config.n_shared_experts
                # 检查当前层是否超过密集层数量(在深层启用MoE)
                and layer_idx >= config.moe_config.n_dense_layer
        )

        # 根据MoE判断结果选择前馈网络类型
        if use_moe:
            # 使用混合专家(MoE)架构
            # MoE通过稀疏激活专家网络来增加模型容量，同时保持计算效率
            # config: MoE配置参数，layer=MLP: 每个专家的基础网络结构
            self.mlp = MoE(config=config, layer=MLP)
        else:
            # 使用标准密集MLP架构
            # 所有输入都通过相同的前馈网络处理
            self.mlp = MLP(config)

    def forward(
            self,
            hidden_states: torch.Tensor,
            position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            past_key_values: Optional[KVCache] = None
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        执行解码器层的前向传播，实现注意力机制和前馈网络的计算。

        该方法实现了标准的Transformer解码器层计算流程，包括预归一化、
        残差连接、注意力计算和前馈网络处理。支持MoE架构并返回辅助损失。

        Args:
            hidden_states (torch.Tensor): 输入隐藏状态
                - 形状: [batch_size, seq_len, hidden_size]
                - 作用: 来自前一层的特征表示，包含序列的语义信息
                - 数值范围: 通常为浮点数，经过归一化处理
            position_embeddings (Optional[Tuple[torch.Tensor, torch.Tensor]]): 旋转位置编码
                - 格式: (cos, sin) 元组
                - cos: 余弦位置编码，形状与query/key相同
                - sin: 正弦位置编码，形状与query/key相同
                - 作用: 为注意力计算注入相对位置信息
                - 默认值: None，在某些情况下可能不需要位置编码
            attention_mask (Optional[torch.Tensor]): 注意力掩码
                - 形状: [batch_size, 1, tgt_seq_len, src_seq_len]
                - 作用: 控制注意力的可见性(因果掩码、填充掩码等)
                - 数值特征: 有效位置为0，无效位置为极小值
                - 默认值: None，表示无特殊注意力限制
            past_key_values (Optional[KVCache]): KV缓存
                - 作用: 存储历史键值对，支持增量推理
                - 应用场景: 文本生成、增量解码等
                - 默认值: None，表示不使用KV缓存

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]: 解码器层的输出和辅助损失
                - hidden_states (torch.Tensor): 处理后的隐藏状态
                    * 形状: [batch_size, seq_len, hidden_size]
                    * 作用: 包含注意力机制和前馈网络处理后的上下文信息
                    * 数值特征: 经过残差连接和归一化的特征表示
                - aux_loss (Optional[torch.Tensor]): MoE辅助损失
                    * 形状: [] (标量张量)
                    * 作用: MoE负载均衡和专家选择正则化损失
                    * 数值特征: 非负数，用于优化专家路由
                    * 默认值: None，当不使用MoE时返回None
        """
        # 步骤1: 保存残差连接的原始输入
        # 残差连接有助于缓解梯度消失问题，支持深层网络训练
        residual = hidden_states

        # 步骤2: 注意力子层计算 (预归一化 + 注意力 + 残差连接)
        # 首先对输入进行层归一化，稳定注意力计算
        # 然后通过多头注意力机制捕获序列内部的依赖关系
        hidden_states = self.attn(
            hidden_states=self.attn_norm(hidden_states),  # 预归一化
            position_embeddings=position_embeddings,  # 位置编码
            attention_mask=attention_mask,  # 注意力掩码
            past_key_values=past_key_values  # KV缓存
        )

        # 应用残差连接：将注意力输出与原始输入相加
        # 这有助于保持梯度流动，防止信息丢失
        hidden_states = hidden_states + residual

        # 步骤3: MLP子层计算 (预归一化 + 前馈网络 + 残差连接)
        # 首先对注意力输出进行层归一化
        normalized_hidden_states = self.mlp_norm(hidden_states)

        # 根据MLP类型(密集或MoE)进行不同的处理
        if isinstance(self.mlp, MoE):
            # MoE架构：返回MLP输出和辅助损失
            # MoE通过路由机制将输入分配给不同的专家网络
            # aux_loss用于负载均衡和专家选择正则化
            mlp_states, aux_loss = self.mlp(normalized_hidden_states)
        else:
            # 密集MLP架构：只返回MLP输出，无辅助损失
            # 所有输入都通过相同的前馈网络处理
            mlp_states = self.mlp(normalized_hidden_states)
            aux_loss = None

        # 应用残差连接：将MLP输出与注意力输出相加
        # 完成解码器层的最终计算
        hidden_states = mlp_states + hidden_states

        # 返回处理后的隐藏状态和可选的MoE辅助损失
        # aux_loss在训练时用于优化专家路由，推理时通常为None
        return hidden_states, aux_loss


class Barranzi2LMForCausalLM(nn.Module):
    """
    大语言模型主类，实现了完整的Transformer架构

    该类是整个大语言模型的核心组件，集成了以下关键功能：
    1. 词嵌入层：将输入token转换为密集向量表示
    2. 旋转位置编码(RoPE)：为token注入位置信息
    3. 多层解码器：通过堆叠DecoderLayer实现深度特征提取
    4. 输出层：将隐藏状态映射到词汇表概率分布
    5. KV缓存：支持增量推理，提高生成效率
    6. MoE支持：可选的混合专家机制，扩展模型容量

    架构特点：
    - 采用仅解码器(Decoder-only)架构，适合自回归文本生成
    - 支持训练和推理两种模式，推理时可使用KV缓存优化
    - 集成RMSNorm归一化，提供稳定的训练过程
    - 灵活的配置系统，支持不同规模和特性的模型变体
    """

    def __init__(self, config: Config):
        """
        初始化大语言模型

        Args:
            config (Config): 模型配置对象，包含模型架构的所有超参数
                           包括词汇表大小、隐藏层维度、层数、注意力头数等
        """
        super().__init__()

        # 旋转位置编码模块，为注意力机制提供位置信息
        # RoPE通过复数旋转的方式将位置信息注入到查询和键向量中
        self.rotary_emb = RotaryEmbedding(config=config)

        # 词嵌入层，将离散的token ID映射为连续的密集向量
        # 形状：(vocab_size, hidden_size)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

        # 解码器层列表，构建模型的深度结构
        # 每个DecoderLayer包含自注意力机制和前馈网络
        # 支持MoE(Mixture of Experts)扩展，可根据配置选择使用标准MLP或MoE
        self.layers = nn.ModuleList(
            [DecoderLayer(config, idx) for idx in range(config.num_hidden_layers)])

        # 输出头归一化层，使用RMSNorm进行最终的隐藏状态归一化
        # RMSNorm相比LayerNorm计算更高效，性能相当
        self.head_norm = RMSNorm(config.hidden_size)

        # 语言模型头，将归一化后的隐藏状态映射到词汇表空间
        # 无偏置设计，减少参数量，提高训练稳定性
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # 应用权重初始化，确保训练开始时的稳定性
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module):
        """
        模型权重初始化方法

        采用正态分布初始化策略，标准差设为0.02，这是Transformer模型的常用设置
        合适的初始化对深度模型的训练收敛至关重要

        Args:
            module (nn.Module): 需要初始化的神经网络模块
        """
        std = 0.02  # 标准差，控制初始化权重的范围

        if isinstance(module, nn.Linear):
            # 线性层权重初始化：均值为0，标准差为0.02的正态分布
            # 这种初始化方式有助于梯度在深层网络中的稳定传播
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                # 偏置项初始化为0，简化模型的学习过程
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            # 嵌入层权重初始化：同样使用正态分布
            # 确保所有token的初始表示都在合理的范围内
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                # 填充token的嵌入向量初始化为0
                # 这些token在计算损失时会被忽略，因此不需要学习有意义的表示
                module.weight.data[module.padding_idx].zero_()

    def forward(
            self,
            input_ids: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            doc_boundary_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            past_key_values: Optional[KVCache] = None,
            use_cache: bool = False,
            logits_to_keep: int = 0,
            **kwargs,
    ) -> Dict[str, any]:
        """
        模型前向传播方法，实现完整的推理流程

        该方法处理输入数据，通过多层Transformer解码器，最终输出预测概率
        支持训练和推理两种模式，推理时可启用KV缓存优化

        Args:
            input_ids (torch.Tensor): 输入token ID张量，形状为(batch_size, seq_len)
                                    包含待处理的文本序列的数字化表示
            attention_mask (torch.Tensor, optional): 注意力掩码，形状为(batch_size, seq_len)
                                                   用于标记padding位置，True表示有效token，False表示padding
                                                   例如：[[True, True, True, False], [True, True, True, True]]
            doc_boundary_mask (torch.Tensor, optional): 文档边界掩码，形状为(batch, 1, seq_len, seq_len)
                                                       用于处理多文档场景，限制注意力范围在文档内部
            position_ids (torch.Tensor, optional): 位置ID张量，形状为(batch_size, seq_len)
                                                  指定每个token的绝对位置，若为None则自动生成
            past_key_values (KVCache, optional): KV缓存对象，用于增量推理
                                               存储之前计算的键值对，避免重复计算
            use_cache (bool): 是否使用KV缓存，默认为False
                            True时启用增量推理模式，提高生成效率
            logits_to_keep (int): 保留最后多少个token的logits，默认为0
                                 用于优化计算，只计算最后几个token的预测概率
            **kwargs: 其他可选参数，用于扩展功能

        Returns:
            Dict[str, any]: 包含以下键的字典：
                - 'logits': 模型输出logits，形状为(batch_size, seq_len, vocab_size)
                           表示每个位置上词汇表中每个词的预测概率
                - 'past_key_values': 更新后的KV缓存，用于下一次增量推理
                - 'aux_loss': MoE辅助损失，如果使用MoE则返回专家负载均衡损失，否则为None
        """
        # 获取输入张量的形状信息
        batch_size, seq_len = input_ids.shape

        # KV缓存初始化：如果启用缓存但未提供缓存对象，则创建新的KVCache
        if use_cache and past_key_values is None:
            past_key_values = KVCache()

        # 获取输入嵌入：将token ID转换为密集向量表示
        # 输出形状：(batch_size, seq_len, hidden_size)
        # 在增量推理模式下，如果使用past_key_values，inputs_embeds形状可能为(1, 1, hidden_size)
        inputs_embeds = self.get_input_embeddings(
            input_ids,
            attention_mask,
            **kwargs
        )

        # 计算历史token数量：从KV缓存中获取已处理的token数量
        # 用于增量推理中的位置编码和注意力掩码计算
        past_seen_tokens = past_key_values.get_seq_len() if past_key_values is not None else 0
        full_seq_len = past_seen_tokens + seq_len  # 总序列长度

        # 位置ID处理：如果未提供位置ID，则自动生成
        # 位置ID从past_seen_tokens开始，确保增量推理中位置编码的连续性
        if position_ids is None:
            position_ids = torch.arange(past_seen_tokens, full_seq_len, device=inputs_embeds.device).unsqueeze(0)

        # 注意力掩码处理：如果未提供掩码，则创建全True掩码
        # 表示所有token都是有效的，没有padding
        if attention_mask is None:
            attention_mask = torch.ones(
                (batch_size, full_seq_len), dtype=torch.bool, device=inputs_embeds.device
            )

        # 计算旋转位置编码：为输入嵌入注入位置信息
        # RoPE通过复数旋转的方式实现位置感知的注意力机制
        position_embeddings = self.rotary_emb(inputs_embeds, position_ids)

        # 准备解码器注意力掩码：整合多种掩码信息
        # 包括因果掩码、填充掩码和文档边界掩码
        # 输出形状：(batch_size, 1, seq_len, full_seq_len)
        causal_mask = prepare_decoder_attention_mask(
            attention_mask=attention_mask,
            doc_boundary_mask=doc_boundary_mask,
            input_shape=(batch_size, seq_len),
            past_key_values_length=past_seen_tokens,
            dtype=inputs_embeds.dtype,
            device=inputs_embeds.device
        )

        # 初始化隐藏状态：从输入嵌入开始逐层处理
        hidden_states = inputs_embeds
        aux_losses = ()  # 用于收集MoE辅助损失

        # 逐层处理：通过所有解码器层进行特征提取
        for layer in self.layers:
            # 每个解码器层执行自注意力和前馈网络计算
            hidden_states, aux_loss = layer(
                hidden_states=hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=causal_mask,
                past_key_values=past_key_values
            )

            # 收集MoE辅助损失：如果该层使用了MoE且有辅助损失
            if aux_loss is not None:
                aux_losses += (aux_loss,)

        # 最终归一化：对最后一层的输出进行RMSNorm归一化
        # 输出形状：(batch_size, seq_len, hidden_size)
        hidden_states = self.head_norm(hidden_states)

        # Logits切片处理：根据logits_to_keep参数选择需要计算的token位置
        # 优化计算效率，避免计算不需要的token预测
        slice_indices = slice(-logits_to_keep, None)
        # 语言模型头：将隐藏状态映射到词汇表空间
        # 输出形状：(batch_size, seq_len, vocab_size)
        head = self.lm_head(hidden_states[:, slice_indices, :])

        # 返回结果：包含logits、KV缓存和辅助损失
        return {
            'logits': head,  # 预测概率分布
            'past_key_values': past_key_values,  # 更新后的KV缓存
            'aux_loss': None if len(aux_losses) == 0 else sum(aux_losses)  # MoE辅助损失总和
        }

    def get_input_embeddings(
            self,
            input_ids: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            **kwargs,
    ):
        """
        获取输入嵌入表示

        将输入的token ID转换为密集的向量表示，这是模型处理的第一步

        Args:
            input_ids (torch.Tensor): 输入token ID张量
            attention_mask (torch.Tensor, optional): 注意力掩码（当前未使用，保留用于扩展）
            **kwargs: 其他可选参数

        Returns:
            torch.Tensor: 词嵌入张量，形状为(batch_size, seq_len, hidden_size)
        """
        return self.embed_tokens(input_ids)